class CliffWalkingEnv:
    def __init__(self, ncol, nrow):
        self.nrow = nrow
        self.ncol = ncol
        self.x = 0  # 记录当前智能体位置的横坐标
        self.y = self.nrow - 1  # 记录当前智能体位置的纵坐标

    def Step(self, action):  # 外部调用这个函数来改变当前位置
        # 4种动作, change[0]:上, change[1]:下, change[2]:左, change[3]:右。坐标系原点(0,0)
        # 定义在左上角
        actions = [[-1, 0], [1, 0], [0, -1], [0, 1]]
        self.x = min(self.nrow - 1, max(0, self.x + actions[action][0]))
        self.y = min(self.ncol - 1, max(0, self.y + actions[action][1]))
        reward = -1
        done = False
        if self.x == self.nrow - 1 and self.y > 0:  # 下一个位置在悬崖或者目标
            done = True
            if self.y != self.ncol - 1:
                reward = -100
        return self.state(), reward, done

    def Reset(self):  # 回归初始状态,坐标轴原点在左上角
        self.x = self.nrow - 1
        self.y = 0
        return self.state()

    def state(self):
        """将坐标转换为状态"""
        return self.x * self.ncol + self.y


def PrintPolicy(agent, env, action_meaning, disaster=[], end=[]):
    for i in range(env.nrow):
        for j in range(env.ncol):
            if (i * env.ncol + j) in disaster:
                print("****", end=" ")
            elif (i * env.ncol + j) in end:
                print("EEEE", end=" ")
            else:
                a = agent.PrintAction(i * env.ncol + j)
                pi_str = ""
                for k in range(len(action_meaning)):
                    pi_str += action_meaning[k] if a[k] > 0 else "o"
                print(pi_str, end=" ")
        print()
